// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <fcntl.h>
#include <sys/resource.h>
#include <sys/time.h>

#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <turbo/flags/flag.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fstream>

#include <pollux/common/base/succinct_printer.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/common/memory/mmap_allocator.h>
#include <pollux/connectors/hive/hive_config.h>
#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/dwio/common/options.h>
#include <pollux/dwio/dwrf/register_dwrf_reader.h>
#include <pollux/dwio/parquet/register_parquet_reader.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/exec/split.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/testing/exec/util/tpch_query_builder.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/parse/type_resolver.h>

TURBO_DECLARE_FLAG(std::string, test_flags_file);
TURBO_DECLARE_FLAG(bool, include_results);
TURBO_DECLARE_FLAG(bool, include_custom_stats);
TURBO_DECLARE_FLAG(std::string, data_format);

namespace kumo::pollux {

struct RunStats {
  std::map<std::string, std::string> flags;
  int64_t micros{0};
  int64_t rawInputBytes{0};
  int64_t userNanos{0};
  int64_t systemNanos{0};
  std::string output;

  std::string toString(bool detail) {
    std::stringstream out;
    out << succinctNanos(micros * 1000) << " "
        << succinctBytes(rawInputBytes / (micros / 1000000.0)) << "/s raw, "
        << succinctNanos(userNanos) << " user " << succinctNanos(systemNanos)
        << " system (" << (100 * (userNanos + systemNanos) / (micros * 1000))
        << "%), flags: ";
    for (auto& pair : flags) {
      out << pair.first << "=" << pair.second << " ";
    }
    out << std::endl << "======" << std::endl;
    if (detail) {
      out << std::endl << output << std::endl;
    }
    return out.str();
  }
};

struct ParameterDim {
  std::string flag;
  std::vector<std::string> values;
};

class QueryBenchmarkBase {
 public:
  virtual ~QueryBenchmarkBase() = default;
  virtual void initialize();
  void shutdown();
  std::pair<std::unique_ptr<exec::TaskCursor>, std::vector<RowVectorPtr>> run(
      const exec::test::TpchPlan& tpchPlan);

  virtual std::vector<std::shared_ptr<connector::ConnectorSplit>> listSplits(
      const std::string& path,
      int32_t numSplitsPerFile,
      const exec::test::TpchPlan& plan);

  static void ensureTaskCompletion(exec::Task* task);

  static bool validateDataFormat(
      const char* flagname,
      const std::string& value);

  static void printResults(
      const std::vector<RowVectorPtr>& results,
      std::ostream& out);

  void readCombinations();

  /// Entry point invoked with different settings to run the benchmark.
  virtual void runMain(std::ostream& out, RunStats& runStats) = 0;

  void runOne(std::ostream& outtt, RunStats& stats);

  void runCombinations(int32_t level);

  void runAllCombinations();

 protected:
  std::unique_ptr<melon::IOThreadPoolExecutor> ioExecutor_;
  std::unique_ptr<melon::IOThreadPoolExecutor> cacheExecutor_;
  std::shared_ptr<memory::MemoryAllocator> allocator_;
  std::shared_ptr<cache::AsyncDataCache> cache_;
  // Parameter combinations to try. Each element specifies a flag and possible
  // values. All permutations are tried.
  std::vector<ParameterDim> parameters_;

  std::vector<RunStats> runStats_;
};
} // namespace kumo::pollux
